package erpc

import (
	"time"
)

type Stream interface {
	Read(b []byte) (n int, err error)
	Write(b []byte) (n int, err error)
	Close() error
}

type limitStream struct {
	conn Stream
	rate int64
}

func NewLimitStream(conn Stream, rate int64) *limitStream {
	return &limitStream{
		conn: conn,
		rate: rate,
	}
}

func (s *limitStream) Read(b []byte) (n int, err error) {
	return s.conn.Read(b)
}

func (s *limitStream) Write(b []byte) (n int, err error) {
	start := time.Now()
	n, err = s.conn.Write(b)
	s.delay(start, n)
	return n, err
}
func (s *limitStream) Close() error {
	return s.conn.Close()
}

func (s *limitStream) delay(start time.Time, n int) {
	if s.rate > 0 {
		expected := time.Duration(int64(time.Second) * int64(n) / s.rate) //second
		if actual := time.Now().Sub(start); expected > actual {
			time.Sleep(expected - actual)
		}
	}
}
